"""
核心坐姿检测类
"""

import cv2
import mediapipe as mp
import numpy as np
import time
from typing import Optional, List, Tuple, Dict
from camera_handler import CameraHandler
from posture_analyzer import PostureAnalyzer, PostureMetrics
from config import POSE_CONFIG, DISPLAY_CONFIG, COLORS
from utils import normalize_landmarks, draw_landmarks, draw_text_with_background, get_grade_color

class PostureDetector:
    """坐姿检测器核心类"""
    
    def __init__(self, camera_id: int = 0):
        """
        初始化坐姿检测器
        
        Args:
            camera_id: 摄像头ID
        """
        self.camera_handler = CameraHandler(camera_id)
        self.posture_analyzer = PostureAnalyzer()
        
        # 初始化MediaPipe Pose
        self.mp_pose = mp.solutions.pose
        self.pose = self.mp_pose.Pose(
            static_image_mode=POSE_CONFIG['static_image_mode'],
            model_complexity=POSE_CONFIG['model_complexity'],
            smooth_landmarks=POSE_CONFIG['smooth_landmarks'],
            enable_segmentation=POSE_CONFIG['enable_segmentation'],
            smooth_segmentation=POSE_CONFIG['smooth_segmentation'],
            min_detection_confidence=POSE_CONFIG['min_detection_confidence'],
            min_tracking_confidence=POSE_CONFIG['min_tracking_confidence']
        )
        
        # MediaPipe绘制工具
        self.mp_drawing = mp.solutions.drawing_utils
        self.mp_drawing_styles = mp.solutions.drawing_styles
        
        # 状态变量
        self.is_running = False
        self.current_metrics = None
        self.frame_count = 0
        
        # 显示配置
        self.display_config = DISPLAY_CONFIG
        self.colors = COLORS
        
        print("坐姿检测器初始化完成")
    
    def detect_posture(self, frame: np.ndarray) -> Dict:
        """
        检测单帧图像中的坐姿
        
        Args:
            frame: 输入图像帧
            
        Returns:
            检测结果字典
        """
        start_time = time.time()
        
        # 转换颜色空间
        rgb_frame = cv2.cvtColor(frame, cv2.COLOR_BGR2RGB)
        
        # MediaPipe姿态检测
        results = self.pose.process(rgb_frame)
        
        # 初始化返回结果
        detection_result = {
            'success': False,
            'landmarks': None,
            'metrics': None,
            'processing_time': 0.0
        }
        
        if results.pose_landmarks:
            # 转换关键点坐标
            landmarks = normalize_landmarks(
                results.pose_landmarks, 
                frame.shape[1], 
                frame.shape[0]
            )
            
            # 分析坐姿
            metrics = self.posture_analyzer.analyze_posture(
                landmarks, 
                frame.shape[1], 
                frame.shape[0]
            )
            
            # 更新检测结果
            detection_result.update({
                'success': True,
                'landmarks': landmarks,
                'metrics': metrics,
                'processing_time': time.time() - start_time
            })
            
            # 更新当前指标
            self.current_metrics = metrics
            self.frame_count += 1
        
        detection_result['processing_time'] = time.time() - start_time
        return detection_result
    
    def start_detection(self) -> bool:
        """
        启动坐姿检测
        
        Returns:
            是否成功启动
        """
        if not self.camera_handler.open_camera():
            print("摄像头初始化失败")
            return False
        
        self.is_running = True
        print("坐姿检测已启动")
        return True
    
    def stop_detection(self):
        """停止坐姿检测"""
        self.is_running = False
        self.camera_handler.close_camera()
        print("坐姿检测已停止")
    
    def reset_detector(self):
        """重置检测器"""
        self.posture_analyzer = PostureAnalyzer()
        self.current_metrics = None
        self.frame_count = 0
        print("检测器已重置")
    
    def get_current_metrics(self) -> Optional[PostureMetrics]:
        """获取当前坐姿指标"""
        return self.current_metrics
    
    def get_detection_stats(self) -> Dict:
        """获取检测统计信息"""
        if not self.current_metrics:
            return {
                'frame_count': self.frame_count,
                'current_score': 0.0,
                'current_grade': 'poor',
                'suggestions': []
            }
        
        return {
            'frame_count': self.frame_count,
            'current_score': self.current_metrics.overall_score,
            'current_grade': self.current_metrics.grade,
            'suggestions': self.current_metrics.suggestions
        }
    
    def is_alert_needed(self, threshold: float = 60.0) -> bool:
        """
        判断是否需要发出警报
        
        Args:
            threshold: 警报阈值
            
        Returns:
            是否需要警报
        """
        if not self.current_metrics:
            return False
        
        return self.current_metrics.overall_score < threshold
    
    def get_posture_summary(self) -> str:
        """获取坐姿摘要"""
        if not self.current_metrics:
            return "未检测到坐姿"
        
        grade_names = {
            'excellent': '优秀',
            'good': '良好',
            'fair': '一般',
            'poor': '较差'
        }
        
        grade_text = grade_names.get(self.current_metrics.grade, '未知')
        return f"坐姿评分: {self.current_metrics.overall_score:.1f}/100 ({grade_text})"
    
    def capture_posture_image(self, filename: str) -> bool:
        """
        捕获当前坐姿图像
        
        Args:
            filename: 保存文件名
            
        Returns:
            是否成功保存
        """
        try:
            success, frame = self.camera_handler.read_frame()
            if success:
                cv2.imwrite(filename, frame)
                return True
        except Exception as e:
            print(f"保存图像失败: {e}")
        
        return False
    
    def get_detailed_analysis(self) -> Dict:
        """获取详细分析报告"""
        if not self.current_metrics:
            return {}
        
        return self.posture_analyzer.get_detailed_analysis(self.current_metrics)
    
    def __del__(self):
        """析构函数"""
        if hasattr(self, 'pose'):
            self.pose.close()
        if hasattr(self, 'camera_handler'):
            self.camera_handler.close_camera()

# 兼容性函数
def create_posture_detector(camera_id: int = 0) -> PostureDetector:
    """
    创建坐姿检测器实例
    
    Args:
        camera_id: 摄像头ID
        
    Returns:
        坐姿检测器实例
    """
    return PostureDetector(camera_id) 